import numpy as np

from ..robot_state import RobotState, LegState
from ..utils.bezier import Bezier


class SwingPlanner:

    def __init__(self) -> None:
        super().__init__()

        self.kp_vx = 0.0
        self.kp_vy = 0.0
        self.kp_wz = 0.0

        self.legLength = 1.0
        self.stepHeight = 0.1
        self.phi = 0.0
        self.tSwing = 0.4
        self.posStart_W = np.zeros(3)
        self.posDes_W = np.zeros(3)
        self.hipPos_W = np.zeros(3)
        self.STPos_W = np.zeros(3)
        self.desV_W = np.zeros(3)
        self.curV_W = np.zeros(3)

        self.desWz_W = 0.0
        self.base_pos = np.zeros(3)
        self.legState = LegState.LST

        self.pDesCur = np.zeros(3)
        self.yawCur = 0.0
        self.theta0 = 0.0
        self.omegaZ_W = 0.0
        self.hip_width = 0.0

    def trajectory(self, phase: float, des1: float, des2: float) -> float:
        Bswpid = Bezier()
        para0 = 5
        para1 = 3
        for i in range(para0):
            Bswpid.push_back(0.0)

        for i in range(para1):
            Bswpid.push_back(1.0)

        output = 0.0
        if self.phi < phase:
            output = des1 * Bswpid.getOut(self.phi / phase)
        else:
            s = Bswpid.getOut((1.0 - self.phi) / (1.0 - phase))
            if s > 0:
                output = des1 * s + des2 * (1.0 - s)
            else:
                output = des2
        return output

    def getSwingPos(self):
        KP = np.zeros((3, 3))
        KP[0, 0] = self.kp_vx
        KP[1, 1] = self.kp_vy

        Rz = np.array([
            [np.cos(self.yawCur), -np.sin(self.yawCur), 0],
            [np.sin(self.yawCur), np.cos(self.yawCur), 0],
            [0, 0, 1],
        ])

        KP = Rz @ KP @ Rz.T

        self.posDes_W[:] = self.hipPos_W + KP @ (self.desV_W - self.curV_W) * (
            -1) + 0.5 * self.tSwing * self.curV_W + self.curV_W * (1 - self.phi) * self.tSwing

        thetaF = self.yawCur + self.theta0 + self.omegaZ_W * (
                1 - self.phi) * self.tSwing + 0.5 * self.omegaZ_W * self.tSwing + self.kp_wz * (
                         self.omegaZ_W - self.desWz_W)
        self.posDes_W[0] += 0.5 * self.hip_width * (np.cos(thetaF) - np.cos(self.yawCur + self.theta0))
        self.posDes_W[1] += 0.5 * self.hip_width * (np.sin(thetaF) - np.sin(self.yawCur + self.theta0))

        xOff_L = 0.02
        yOff_L = 0.0
        zOff_W = -0.032

        self.posDes_W[2] = self.base_pos[2] - self.legLength + zOff_W

        xOff_W = 0.0
        yOff_W = 0.0
        if self.legState == LegState.LST:
            xOff_W = np.cos(self.yawCur) * xOff_L - np.sin(self.yawCur) * yOff_L
            yOff_W = np.sin(self.yawCur) * xOff_L + np.cos(self.yawCur) * yOff_L
        elif self.legState == LegState.RST:
            xOff_W = np.cos(self.yawCur) * xOff_L - np.sin(self.yawCur) * (-yOff_L)
            yOff_W = np.sin(self.yawCur) * xOff_L + np.cos(self.yawCur) * (-yOff_L)

        self.posDes_W[0] += xOff_W
        self.posDes_W[1] += yOff_W

        if self.phi < 1.0:
            self.pDesCur[0] = self.posStart_W[0] + (self.posDes_W[0] - self.posStart_W[0]) / (2 * np.pi) * (
                    2 * np.pi * self.phi - np.sin(2 * np.pi * self.phi))
            self.pDesCur[1] = self.posStart_W[1] + (self.posDes_W[1] - self.posStart_W[1]) / (2 * np.pi) * (
                    2 * np.pi * self.phi - np.sin(2 * np.pi * self.phi))

        self.pDesCur[2] = self.posStart_W[2] + self.trajectory(0.2, self.stepHeight,
                                                               self.posDes_W[2] - self.posStart_W[2])

    def read_state(self, robot_state: RobotState):
        self.posStart_W[:] = robot_state.swingStartPos_W
        self.desV_W[:] = robot_state.js_vel_des
        self.desWz_W = robot_state.js_vel_des[2]
        self.curV_W[:] = robot_state.dq[:3]
        self.phi = robot_state.phi
        self.hipPos_W[:] = robot_state.posHip_W
        self.STPos_W[:] = robot_state.posST_W
        self.base_pos[:] = robot_state.base_pos
        self.tSwing = robot_state.tSwing
        self.theta0 = robot_state.theta0
        self.yawCur = robot_state.rpy[2]
        self.omegaZ_W = robot_state.base_omega_W[2]
        self.hip_width = robot_state.width_hips
        self.legState = robot_state.legState

    def write_state(self, robot_state: RobotState):
        robot_state.swingDesPosCur_W[:3] = self.pDesCur[:3]
        robot_state.swingDesPosFinal_W[:3] = self.posDes_W[:3]
        robot_state.swing_fe_rpy_des_W[:3] = [0.0, 0.0, robot_state.base_rpy_des[2]]
        robot_state.swing_fe_pos_des_W[:3] = self.pDesCur[:3]
